import argparse


def parsers():
    parser = argparse.ArgumentParser()
    parser.add_argument('--num_classes', type = int, default = 7)
    parser.add_argument('--epochs', type = int, default = 2)
    parser.add_argument('--batch_size', type = int, default = 8)
    parser.add_argument('--lr', type = float, default = 0.001)
    parser.add_argument('--lrf', type = float, default = 0.01)
    parser.add_argument('--weight_decay', type = float, default = 5e-5)
    parser.add_argument('--droput', type = float, default = 0.5)

    # 数据集所在根目录
    parser.add_argument('--data_path', type = str,
                        default = r"./image_data")
    parser.add_argument('--model_name', default = 'Beit', help = 'model name')

    # 预训练权重路径，如果不想载入就设置为空字符
    parser.add_argument('--weights', type = str, default = './beit/beit_base_patch16_224_pt22k.pth',
                        help = 'initial weights path')

    # 是否冻结权重
    parser.add_argument('--freeze_layers', type = bool, default = False)
    parser.add_argument('--device', default = 'cuda:0', help = 'device id (i.e. 0 or 0,1 or cpu)')

    # 保存训练的模型
    parser.add_argument('--best_model', type = str, default = './beit_model/best_model.pth',
                        help = 'best_model_save_pth')
    args = parser.parse_args()
    return args
